import os  # 用于访问环境变量和操作系统功能

from langchain_community.embeddings import DashScopeEmbeddings
from langchain_community.vectorstores import Chroma
from langchain_core.documents import Document
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables import RunnableLambda, RunnablePassthrough
from langchain_openai import ChatOpenAI  # 用于调用OpenAI兼容接口的模型（适配阿里云百炼）
from pydantic import SecretStr  # 用于安全存储和传递API密钥（避免明文暴露）

api_key = os.getenv("DASHSCOPE_API_KEY")
# 检查API密钥是否存在，若不存在则抛出错误提示
if not api_key:
    raise ValueError("请设置环境变量DASHSCOPE_API_KEY（阿里云百炼API-KEY）")
# 1. 创建大语言模型实例（适配阿里云百炼的DeepSeek模型）
model = ChatOpenAI(
    model="deepseek-r1-distill-qwen-1.5b",  # 模型名称，需与阿里云百炼平台支持的模型名一致
    temperature=0,  # 生成内容的随机性（0表示确定性输出，1表示最大随机性）
    max_tokens=None,  # 最大生成 tokens 数（None表示使用模型默认值）
    timeout=None,  # 超时时间（None表示无超时限制）
    max_retries=2,  # 调用失败时的最大重试次数
    api_key=SecretStr(api_key),  # 以安全字符串形式传递API密钥
    base_url="https://dashscope.aliyuncs.com/compatible-mode/v1",  # 阿里云百炼的OpenAI兼容接口地址
)

# 准备数据
documents = [
    Document(
        page_content="狗是伟大的伴侣，以其忠诚和友好而闻名。",
        metadata={'source': '哺乳动物文档'},
    ),
    Document(
        page_content="猫是独立的宠物，通常喜欢自己的空间。",
        metadata={'source': '哺乳动物文档'},
    ),
    Document(
        page_content="金鱼是初学者的流行宠物，需要相对简单的护理",
        metadata={'source': '鱼类宠物文档'},
    ),
]

# 实例化一个向量数据库空间
vector_stores = Chroma.from_documents(
    documents,
    embedding=DashScopeEmbeddings(
        model="text-embedding-v4",
    ))  # embedding向量数据的方式

# 检索器  bind(k=1)返回相似度最高的一个
retriever = RunnableLambda(vector_stores.similarity_search).bind(k=1)

# 2.创建prompt提升模板
message = """
使用提供的上下文仅回答这个问题。
{question}
上下文:
{context}
"""

prompt_template = ChatPromptTemplate.from_messages([('human', message)])

# 3.组成链路
# RunnablePassthrough允许我们将用户的问题 后面再传递给 prompt或model
chain = {'question': RunnablePassthrough(), 'context': retriever} | prompt_template | model

# 4.调用
res = chain.invoke('请介绍猫')

print(res.content)
